package com.wing.tools.video;

import android.graphics.ImageFormat;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.os.Build;
import android.support.annotation.RequiresApi;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.List;

import static android.media.MediaCodec.BUFFER_FLAG_CODEC_CONFIG;
import static android.media.MediaCodec.BUFFER_FLAG_KEY_FRAME;

/**
 * Created by wejust on 2017/12/6.
 */

public class AvcEncoder {
    public interface  IEncodeCallback{
        void OnData(byte[] data);
    }

    private int TIMEOUT_USEC = 12000;

    private MediaCodec mediaCodec;
    private int m_width;
    private int m_height;
    private byte[] yuv420 = null;
    private byte[] m_info = null;
    private int m_keyFrameInterval;
    private int m_imageFormat = 17;
    private int m_framerate;
    private int m_inputIndex = 0;

    public MediaFormat mediaFormat;
    private int timer = -1;
    private long startTime = 0;
    private long recordStartTime = 0;

    @RequiresApi(api = Build.VERSION_CODES.JELLY_BEAN)
    public AvcEncoder(int width, int height, int framerate, int bitrate, int kfInterval, int imageFormat){

        m_width  = width;
        m_height = height;
        m_keyFrameInterval = kfInterval;
        yuv420 = new byte[width*height*3/2];
        m_imageFormat = imageFormat;
        m_framerate = framerate;

        try {
            mediaCodec = MediaCodec.createEncoderByType("video/avc");
        } catch (IOException e) {
            e.printStackTrace();
        }

        mediaFormat = MediaFormat.createVideoFormat("video/avc", width, height);
        mediaFormat.setInteger(MediaFormat.KEY_CHANNEL_COUNT, 1);
        mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, bitrate);
        mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, framerate);
        switch(imageFormat ){
            case ImageFormat.NV21:
                mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420SemiPlanar);
                break;
            case ImageFormat.YV12:
                mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar);
                break;
            default:
                break;
        }
        mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Planar);
        mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 0); //关键帧间隔时间 单位s

        mediaCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
    }

    @RequiresApi(api = Build.VERSION_CODES.JELLY_BEAN)
    public  void start(){
        mediaCodec.start();

        m_inputIndex = 0;

        timer = -1;
        startTime = System.currentTimeMillis();
        recordStartTime = startTime;
    }

    @RequiresApi(api = Build.VERSION_CODES.JELLY_BEAN)
    public void close() {
        try {
            mediaCodec.stop();
            mediaCodec.release();

            m_inputIndex = 0;
        } catch (Exception e){
            e.printStackTrace();
        }
    }

    /**
     * Generates the presentation time for frame N, in microseconds.
     */
    private long computePresentationTime(long frameIndex) {
        //return 132 + frameIndex * 1000000 / m_framerate;
        return System.currentTimeMillis() - recordStartTime;
    }

    @RequiresApi(api = Build.VERSION_CODES.JELLY_BEAN)
    public int offerEncoder(byte[] input, IEncodeCallback cb)//byte[] output)
    {
        int pos = 0;
        if(m_imageFormat == ImageFormat.YV12) {
            swapYV12ToI420(input, yuv420, m_width, m_height);
        }else{
            swapNV21ToNV12(input, yuv420, m_width, m_height);
        }
        try {
            ByteBuffer[] inputBuffers = mediaCodec.getInputBuffers();
            ByteBuffer[] outputBuffers = mediaCodec.getOutputBuffers();
            int inputBufferIndex = mediaCodec.dequeueInputBuffer(-1);
            if (inputBufferIndex >= 0)
            {
                ByteBuffer inputBuffer = inputBuffers[inputBufferIndex];
                inputBuffer.clear();
                inputBuffer.put(yuv420);

                long dts = computePresentationTime(m_inputIndex);
                mediaCodec.queueInputBuffer(inputBufferIndex, 0, yuv420.length, dts, 0);

                m_inputIndex++;
            }

            MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
            int outputBufferIndex = mediaCodec.dequeueOutputBuffer(bufferInfo,TIMEOUT_USEC);

            while (outputBufferIndex >= 0)
            {
                ByteBuffer outputBuffer = outputBuffers[outputBufferIndex];
                byte[] outData = new byte[bufferInfo.size];
                outputBuffer.get(outData);

//                if(m_info != null)
//                {
//                    System.arraycopy(outData, 0,  output, pos, outData.length);
//                    pos += outData.length;
//                }
//                else //保存pps sps 只有开始时 第一个帧里有， 保存起来后面用
//                {
//                    ByteBuffer spsPpsBuffer = ByteBuffer.wrap(outData);
//                    if (spsPpsBuffer.getInt() == 0x00000001)
//                    {
//                        m_info = new byte[outData.length];
//                        System.arraycopy(outData, 0, m_info, 0, outData.length);
//                    }
//                    else
//                    {
//                        return -1;
//                    }
//                }


                if(bufferInfo.flags == BUFFER_FLAG_CODEC_CONFIG){
                    m_info = new byte[outData.length];
                    System.arraycopy(outData, 0, m_info, 0, outData.length);
                }else if(timer < 0 || timer > m_keyFrameInterval * 1000){ // if((bufferInfo.flags & MediaCodec.BUFFER_FLAG_KEY_FRAME) != 0){
                    int nalu_len = 0;
                    if(outData[0] == 0 && outData[1] ==0 && outData[2] == 0 && outData[3] == 1){
                        nalu_len = 5;
                    }
                    byte[] keyframe = new byte[bufferInfo.size + m_info.length - nalu_len];
                    System.arraycopy(m_info, 0, keyframe, 0, m_info.length);
                    //把编码后的视频帧从编码器输出缓冲区中拷贝出来
                    System.arraycopy(outData, nalu_len, keyframe, m_info.length, outData.length - nalu_len);

                    cb.OnData(keyframe);

                    timer = 0;
                    startTime = System.currentTimeMillis();
                }else{
                    cb.OnData(outData);
                    timer += (System.currentTimeMillis() - startTime);
                }


                mediaCodec.releaseOutputBuffer(outputBufferIndex, false);
                outputBufferIndex = mediaCodec.dequeueOutputBuffer(bufferInfo, TIMEOUT_USEC);
            }

//            if(output[4] == 0x65) //key frame   编码器生成关键帧时只有 00 00 00 01 65 没有pps sps， 要加上
//            {
//                System.arraycopy(output, 0,  yuv420, 0, pos);
//                System.arraycopy(m_info, 0,  output, 0, m_info.length);
//                System.arraycopy(yuv420, 0,  output, m_info.length, pos);
//                pos += m_info.length;
//            }

        } catch (Throwable t) {
            t.printStackTrace();
        }

        return pos;
    }
    //yv12 转 yuv420p  yvu -> yuv
    private void swapYV12ToI420(byte[] yv12bytes, byte[] i420bytes, int width, int height)
    {
        System.arraycopy(yv12bytes, 0, i420bytes, 0,width*height);
        System.arraycopy(yv12bytes, width*height+width*height/4, i420bytes, width*height,width*height/4);
        System.arraycopy(yv12bytes, width*height, i420bytes, width*height+width*height/4,width*height/4);
    }

    //把待编码的视频帧转换为YUV420格式
    private void swapNV21ToNV12(byte[] nv21,byte[] nv12,int width,int height){
        if(nv21 == null || nv12 == null)return;
        int framesize = width*height;
        int i = 0,j = 0;
        System.arraycopy(nv21, 0, nv12, 0, framesize);
        for(i = 0; i < framesize; i++){
            nv12[i] = nv21[i];
        }
        for (j = 0; j < framesize/2; j+=2)
        {
            nv12[framesize + j-1] = nv21[j+framesize];
        }
        for (j = 0; j < framesize/2; j+=2)
        {
            nv12[framesize + j] = nv21[j+framesize-1];
        }
    }
}
